import torch.nn as nn


class PreHead(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(PreHead, self).__init__()

        convs = []
        convs.append(conv3x3_bn_relu(in_channels[-1], out_channels))
        # convs.append(conv3x3_bn_relu(channels, channels))
        self.convs = nn.Sequential(*convs)

    def forward(self, x):
        output = self.convs(x[-1])
        return output


def conv3x3_bn_relu(in_planes, out_planes, stride=1):
    "3x3 convolution + BN + relu"
    return nn.Sequential(
            nn.Conv2d(in_planes, out_planes, kernel_size=3,
                      stride=stride, padding=1, bias=False),
            nn.BatchNorm2d(out_planes),
            nn.ReLU(inplace=True),
            )
